"""
DJI Thermal SDK处理器模块
负责R-JPEG文件的读取、温度数据提取和分析
"""

import os
import sys
import numpy as np
import cv2
import ctypes
from typing import Optional, Tuple, Dict, Any
import logging
from pathlib import Path


class TSDKProcessor:
    """DJI Thermal SDK处理器"""
    
    def __init__(self):
        self.sdk_initialized = False
        self.sdk_lib = None
        self.tsdk_path = None
        self.logger = logging.getLogger(__name__)
        
        # 温度数据缓存
        self.current_temp_matrix = None
        self.current_image = None
        self.current_metadata = {}
        
    def initialize_sdk(self, sdk_path: str = None) -> bool:
        """
        初始化DJI Thermal SDK
        
        Args:
            sdk_path: TSDK安装路径，如果为None则自动检测
            
        Returns:
            bool: 初始化是否成功
        """
        try:
            # 自动检测或使用指定的SDK路径
            if sdk_path is None:
                sdk_path = self._detect_sdk_path()
            
            if not sdk_path or not os.path.exists(sdk_path):
                self.logger.warning("TSDK未找到，将使用模拟模式")
                return self._initialize_mock_mode()
            
            self.tsdk_path = sdk_path
            
            # 加载TSDK动态库
            if sys.platform == "win32":
                lib_name = "libdirp.dll"
            else:
                lib_name = "libdirp.so"
                
            lib_path = os.path.join(sdk_path, "lib", lib_name)
            
            if os.path.exists(lib_path):
                # 添加DLL搜索路径
                if hasattr(os, 'add_dll_directory'):
                    os.add_dll_directory(os.path.abspath(os.path.dirname(lib_path)))
                
                self.sdk_lib = ctypes.CDLL(lib_path)
                self._setup_api_functions()
                self.sdk_initialized = True
                # 真实SDK模式，不设置mock_mode
                self.logger.info(f"真实TSDK初始化成功: {lib_path}")
                return True
            else:
                self.logger.warning(f"TSDK库文件未找到: {lib_path}")
                return self._initialize_mock_mode()
                
        except Exception as e:
            self.logger.error(f"TSDK初始化失败: {e}")
            return self._initialize_mock_mode()
    
    def _detect_sdk_path(self) -> Optional[str]:
        """自动检测TSDK安装路径"""
        possible_paths = [
            "tsdk",
            "../tsdk",
            "../../tsdk",
            "C:/Program Files/DJI/TSDK",
            "/usr/local/tsdk",
            "/opt/tsdk"
        ]
        
        for path in possible_paths:
            full_path = os.path.abspath(path)
            if os.path.exists(full_path):
                return full_path
        
        return None
    
    def _initialize_mock_mode(self) -> bool:
        """初始化模拟模式（用于开发和测试）"""
        self.sdk_initialized = True
        self.mock_mode = True
        self.logger.info("TSDK模拟模式已启用")
        return True
    
    def _setup_api_functions(self):
        """设置API函数签名 - 基于真实TSDK API"""
        if not self.sdk_lib:
            return
            
        try:
            # 基于dirp_api.h的实际函数签名
            
            # dirp_create_from_rjpeg(const uint8_t *data, int32_t size, DIRP_HANDLE *ph)
            self.sdk_lib.dirp_create_from_rjpeg.argtypes = [
                ctypes.POINTER(ctypes.c_ubyte),  # const uint8_t *data
                ctypes.c_int32,                  # int32_t size  
                ctypes.POINTER(ctypes.c_void_p)  # DIRP_HANDLE *ph
            ]
            self.sdk_lib.dirp_create_from_rjpeg.restype = ctypes.c_int32
            
            # dirp_get_rjpeg_resolution(DIRP_HANDLE h, dirp_resolution_t *resolution)
            self.sdk_lib.dirp_get_rjpeg_resolution.argtypes = [
                ctypes.c_void_p,                 # DIRP_HANDLE h
                ctypes.POINTER(ctypes.c_int32 * 2)  # dirp_resolution_t* (width, height)
            ]
            self.sdk_lib.dirp_get_rjpeg_resolution.restype = ctypes.c_int32
            
            # dirp_measure(DIRP_HANDLE h, int16_t *temp_image, int32_t size)
            self.sdk_lib.dirp_measure.argtypes = [
                ctypes.c_void_p,                 # DIRP_HANDLE h
                ctypes.POINTER(ctypes.c_int16),  # int16_t *temp_image
                ctypes.c_int32                   # int32_t size
            ]
            self.sdk_lib.dirp_measure.restype = ctypes.c_int32
            
            # dirp_measure_ex(DIRP_HANDLE h, float *temp_image, int32_t size) - 浮点版本
            self.sdk_lib.dirp_measure_ex.argtypes = [
                ctypes.c_void_p,                 # DIRP_HANDLE h
                ctypes.POINTER(ctypes.c_float),  # float *temp_image
                ctypes.c_int32                   # int32_t size
            ]
            self.sdk_lib.dirp_measure_ex.restype = ctypes.c_int32
            
            # dirp_destroy(DIRP_HANDLE h)
            self.sdk_lib.dirp_destroy.argtypes = [ctypes.c_void_p]  # DIRP_HANDLE h
            self.sdk_lib.dirp_destroy.restype = ctypes.c_int32
            
            self.logger.info("TSDK API函数签名设置完成")
            
        except AttributeError as e:
            self.logger.warning(f"部分API函数不可用: {e}")
    
    def process_rjpeg_file(self, file_path: str) -> Dict[str, Any]:
        """
        处理R-JPEG文件
        
        Args:
            file_path: R-JPEG文件路径
            
        Returns:
            Dict: 包含温度矩阵、图像和元数据的字典
        """
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件不存在: {file_path}")
        
        if not self.sdk_initialized:
            raise RuntimeError("TSDK未初始化")
        
        # 检查文件是否为R-JPEG格式
        if not self._is_rjpeg_file(file_path):
            raise ValueError(f"不是有效的R-JPEG文件: {file_path}")
        
        if hasattr(self, 'mock_mode') and self.mock_mode:
            return self._process_rjpeg_mock(file_path)
        
        return self._process_rjpeg_real(file_path)
    
    def _is_rjpeg_file(self, file_path: str) -> bool:
        """检查文件是否为R-JPEG格式"""
        try:
            # 检查文件头标识
            with open(file_path, 'rb') as f:
                header = f.read(12)
                # R-JPEG文件通常在JPEG头部后有特殊标识
                if header[:2] == b'\xff\xd8':  # JPEG文件头
                    # 进一步检查是否包含热成像数据段
                    f.seek(0)
                    content = f.read(1024)  # 读取前1KB检查
                    return b'FLIR' in content or b'DJI' in content or file_path.lower().endswith(('.r.jpg', '.rjpeg'))
            return False
        except Exception:
            return False
    
    def _process_rjpeg_real(self, file_path: str) -> Dict[str, Any]:
        """使用真实TSDK处理R-JPEG文件"""
        obj_ptr = ctypes.c_void_p()
        
        try:
            # 读取R-JPEG文件数据
            with open(file_path, 'rb') as f:
                file_data = f.read()
            
            # 转换为ctypes数组
            data_size = len(file_data)
            data_buffer = (ctypes.c_ubyte * data_size).from_buffer_copy(file_data)
            
            # 创建TSDK对象
            ret = self.sdk_lib.dirp_create_from_rjpeg(
                data_buffer,    # const uint8_t *data
                data_size,      # int32_t size
                ctypes.byref(obj_ptr)  # DIRP_HANDLE *ph
            )
            
            if ret != 0:
                raise RuntimeError(f"创建TSDK对象失败, 错误代码: {ret}")
            
            # 获取图像尺寸 - 使用正确的结构体
            resolution = (ctypes.c_int32 * 2)()  # [width, height]
            ret = self.sdk_lib.dirp_get_rjpeg_resolution(
                obj_ptr, 
                resolution
            )
            
            if ret != 0:
                raise RuntimeError(f"获取图像尺寸失败, 错误代码: {ret}")
            
            width, height = resolution[0], resolution[1]
            pixel_count = width * height
            temp_size = pixel_count * 4  # 每个float占4字节
            
            # 分配温度数据缓冲区 - 使用浮点版本
            temp_buffer = (ctypes.c_float * pixel_count)()
            
            # 提取温度数据 - 使用dirp_measure_ex获取浮点数据
            ret = self.sdk_lib.dirp_measure_ex(obj_ptr, temp_buffer, temp_size)
            
            if ret != 0:
                raise RuntimeError(f"提取温度数据失败, 错误代码: {ret}")
            
            # 转换为numpy数组
            temp_array = np.array(temp_buffer).reshape((height, width))
            
            # 读取可见光图像（如果可能的话，否则使用温度图像）
            try:
                # 使用cv2.imdecode避免中文路径问题
                with open(file_path, 'rb') as f:
                    image_data = f.read()
                image_array = np.frombuffer(image_data, dtype=np.uint8)
                visible_image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
                
                if visible_image is None:
                    self.logger.warning(f"无法读取可见光图像，使用温度数据生成伪彩色图像")
                    # 如果无法读取可见光，创建伪彩色图像
                    visible_image = self._create_pseudocolor_image(temp_array)
            except Exception as e:
                self.logger.warning(f"读取可见光图像失败: {e}，使用温度数据生成伪彩色图像")
                visible_image = self._create_pseudocolor_image(temp_array)
            
            # 检测可见光图像与热成像数据的分辨率差异
            vis_height, vis_width = visible_image.shape[:2]
            temp_height, temp_width = temp_array.shape
            
            # 构建结果
            result = {
                'temperature_matrix': temp_array,
                'visible_image': visible_image,
                'width': width,
                'height': height,
                'file_path': file_path,
                'processing_mode': 'tsdk_real',
                'min_temp': float(np.min(temp_array)),
                'max_temp': float(np.max(temp_array)),
                'mean_temp': float(np.mean(temp_array)),
                'thermal_resolution': {'width': temp_width, 'height': temp_height},
                'visible_resolution': {'width': vis_width, 'height': vis_height},
                'needs_coordinate_mapping': (vis_width != temp_width or vis_height != temp_height),
                'scale_factors': {
                    'x': temp_width / vis_width if vis_width > 0 else 1.0,
                    'y': temp_height / vis_height if vis_height > 0 else 1.0
                }
            }
            
            # 缓存当前数据
            self.current_temp_matrix = temp_array
            self.current_image = visible_image
            self.current_metadata = result
            
            self.logger.info(f"真实TSDK处理完成: 热成像{temp_width}x{temp_height}, 可见光{vis_width}x{vis_height}, 温度范围: {result['min_temp']:.1f}°C - {result['max_temp']:.1f}°C")
            if result['needs_coordinate_mapping']:
                self.logger.info(f"需要坐标映射: 缩放比例 x={result['scale_factors']['x']:.3f}, y={result['scale_factors']['y']:.3f}")
            
            return result
            
        finally:
            # 清理TSDK对象
            if obj_ptr.value:
                self.sdk_lib.dirp_destroy(obj_ptr)
    
    def _process_rjpeg_mock(self, file_path: str) -> Dict[str, Any]:
        """模拟模式处理R-JPEG文件"""
        try:
            # 读取图像 - 使用cv2.imdecode避免中文路径问题
            try:
                # 方法1：直接读取
                image = cv2.imread(file_path, cv2.IMREAD_COLOR)
                if image is None:
                    # 方法2：使用imdecode处理中文路径和特殊格式
                    with open(file_path, 'rb') as f:
                        image_data = f.read()
                    image_array = np.frombuffer(image_data, dtype=np.uint8)
                    image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
                    
                if image is None:
                    raise ValueError(f"无法读取图像文件: {file_path}")
                    
            except Exception as e:
                self.logger.error(f"模拟模式图像读取失败: {e}")
                raise ValueError(f"无法读取图像: {file_path} - {str(e)}")
            
            height, width = image.shape[:2]
            
            # 生成模拟温度数据
            temp_matrix = self._generate_mock_temperature_data(image)
            
            result = {
                'temperature_matrix': temp_matrix,
                'visible_image': image,
                'width': width,
                'height': height,
                'file_path': file_path,
                'processing_mode': 'mock',
                'min_temp': float(np.min(temp_matrix)),
                'max_temp': float(np.max(temp_matrix)),
                'mean_temp': float(np.mean(temp_matrix))
            }
            
            # 缓存数据
            self.current_temp_matrix = temp_matrix
            self.current_image = image
            self.current_metadata = result
            
            return result
            
        except Exception as e:
            raise RuntimeError(f"模拟模式处理失败: {e}")
    
    def _generate_mock_temperature_data(self, image: np.ndarray) -> np.ndarray:
        """基于图像生成模拟温度数据"""
        height, width = image.shape[:2]
        
        # 转换为灰度图
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image
        
        # 基于亮度生成温度数据
        # 假设温度范围20-80°C，亮度越高温度越高
        temp_matrix = 20 + (gray.astype(np.float32) / 255.0) * 60
        
        # 添加一些随机噪声使其更真实
        noise = np.random.normal(0, 2, temp_matrix.shape)
        temp_matrix += noise
        
        # 添加一些热点模拟
        center_y, center_x = height // 2, width // 2
        y, x = np.ogrid[:height, :width]
        
        # 中心热点
        center_dist = np.sqrt((x - center_x) ** 2 + (y - center_y) ** 2)
        center_heat = np.exp(-center_dist / (width * 0.1)) * 10
        
        temp_matrix += center_heat
        
        return temp_matrix
    
    def _create_pseudocolor_image(self, temp_matrix: np.ndarray) -> np.ndarray:
        """基于温度数据创建伪彩色图像"""
        # 标准化温度数据到0-255范围
        temp_normalized = ((temp_matrix - np.min(temp_matrix)) / 
                          (np.max(temp_matrix) - np.min(temp_matrix) + 1e-6) * 255).astype(np.uint8)
        
        # 使用热力图颜色映射
        pseudocolor = cv2.applyColorMap(temp_normalized, cv2.COLORMAP_JET)
        
        return pseudocolor
    
    def get_temperature_at_point(self, x: int, y: int) -> float:
        """获取指定点的温度"""
        if self.current_temp_matrix is None:
            raise RuntimeError("没有可用的温度数据")
        
        height, width = self.current_temp_matrix.shape
        
        if x < 0 or x >= width or y < 0 or y >= height:
            raise ValueError(f"坐标超出范围: ({x}, {y}), 图像尺寸: {width}x{height}")
        
        return float(self.current_temp_matrix[y, x])
    
    def get_region_temperature_stats(self, x1: int, y1: int, x2: int, y2: int) -> Dict[str, Any]:
        """获取区域温度统计"""
        if self.current_temp_matrix is None:
            raise RuntimeError("没有可用的温度数据")
        
        # 确保坐标顺序正确
        x1, x2 = min(x1, x2), max(x1, x2)
        y1, y2 = min(y1, y2), max(y1, y2)
        
        # 提取区域
        region = self.current_temp_matrix[y1:y2, x1:x2]
        
        if region.size == 0:
            raise ValueError("区域为空")
        
        # 计算统计信息
        stats = {
            'min_temp': float(np.min(region)),
            'max_temp': float(np.max(region)),
            'mean_temp': float(np.mean(region)),
            'std_temp': float(np.std(region)),
            'median_temp': float(np.median(region))
        }
        
        # 找到极值位置（相对于区域）
        min_pos = np.unravel_index(np.argmin(region), region.shape)
        max_pos = np.unravel_index(np.argmax(region), region.shape)
        
        # 转换为绝对坐标
        stats['min_position'] = {'x': x1 + min_pos[1], 'y': y1 + min_pos[0]}
        stats['max_position'] = {'x': x1 + max_pos[1], 'y': y1 + max_pos[0]}
        
        return stats
    
    def export_temperature_data(self, output_path: str, format: str = 'csv') -> bool:
        """导出温度数据"""
        if self.current_temp_matrix is None:
            raise RuntimeError("没有可用的温度数据")
        
        try:
            if format.lower() == 'csv':
                np.savetxt(output_path, self.current_temp_matrix, delimiter=',', fmt='%.2f')
            elif format.lower() == 'npy':
                np.save(output_path, self.current_temp_matrix)
            else:
                raise ValueError(f"不支持的导出格式: {format}")
            
            return True
            
        except Exception as e:
            self.logger.error(f"导出温度数据失败: {e}")
            return False
    
    def cleanup(self):
        """清理资源"""
        self.current_temp_matrix = None
        self.current_image = None
        self.current_metadata = {}
        
        if self.sdk_lib:
            # 这里可以添加SDK的清理函数调用
            pass


# 全局TSDK处理器实例
_tsdk_processor = None

def get_tsdk_processor() -> TSDKProcessor:
    """获取全局TSDK处理器实例"""
    global _tsdk_processor
    if _tsdk_processor is None:
        _tsdk_processor = TSDKProcessor()
    return _tsdk_processor